#include <omnetpp.h>
#include "ComputingMessages_m.h"  // 确保正确包含消息头文件
#include <memory>

// 手动实现 make_unique
template<typename T, typename... Args>
std::unique_ptr<T> make_unique(Args&&... args) {
    return std::unique_ptr<T>(new T(std::forward<Args>(args)...));
}

using namespace omnetpp;
using namespace GPUsharing;  // 假设使用了 GPUsharing 命名空间

class TransferServer : public cSimpleModule
{
  protected:
    std::map<int, std::pair<int, int>> requesterInfo;  // 存储请求者信息 <请求者 ID, <所需算力, 任务 ID>>
    std::string userNodeModuleName; // 用于存储用户节点子模块名称
    virtual void handleMessage(cMessage *msg) override;
    int findOutputGateIndex(int nodeId);
    virtual void initialize() override; // 添加初始化函数声明
};

Define_Module(TransferServer);

void TransferServer::initialize() {
    userNodeModuleName = par("userNodeModuleName").stdstringValue(); // 从配置获取用户节点模块名称
}

int TransferServer::findOutputGateIndex(int nodeId) {
    for (int i = 0; i < gateSize("out"); i++) {
        if (getParentModule()->getSubmodule(userNodeModuleName.c_str(), i)->par("nodeId").intValue() == nodeId) {
            return i;
        }
    }
    return -1;
}

void TransferServer::handleMessage(cMessage *msg)
{
    if (auto request = dynamic_cast<RequestComputingPower*>(msg)) {
        // 处理请求消息
        for (int i = 0; i < gateSize("out"); i++) {
            auto newRequest = make_unique<RequestComputingPower>(*request);
            send(newRequest.release(), "out", i);
        }
        requesterInfo[request->getRequesterId()] = std::make_pair(request->getRequiredPower(), 1);
        delete request;
    } else if (auto response = dynamic_cast<ResponseComputingPower*>(msg)) {
        // 处理响应消息
        for (const auto& info : requesterInfo) {
            int requesterId = info.first;
            int requiredPower = info.second.first;
            if (response->getHasEnoughPower() && response->getAvailablePower() >= requiredPower) {
                auto newResponse = make_unique<ResponseComputingPower>("Response");
                newResponse->setResponderId(response->getResponderId());
                newResponse->setHasEnoughPower(true);
                newResponse->setAvailablePower(response->getAvailablePower());
                int gateIndex = findOutputGateIndex(requesterId);
                if (gateIndex != -1) {
                    send(newResponse.release(), "out", gateIndex);
                }
            }
        }
        delete response;
    } else if (auto assignment = dynamic_cast<TaskAssignment*>(msg)) {
        // 处理任务分配消息
        int gateIndex = findOutputGateIndex(assignment->getResponderId());
        if (gateIndex != -1) {
            auto newAssignment = make_unique<TaskAssignment>(*assignment);
            send(newAssignment.release(), "out", gateIndex);
        }
        delete assignment;
    } else if (auto result = dynamic_cast<TaskResult*>(msg)) {
        // 处理任务结果消息
        int gateIndex = findOutputGateIndex(result->getRequesterId());
        if (gateIndex != -1) {
            auto newResult = make_unique<TaskResult>(*result);
            send(newResult.release(), "out", gateIndex);
        }
        delete result;
    } else {
        EV_ERROR << "Received unknown message type: " << msg->getClassName() << endl;
        delete msg;
    }
}
